#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import open3d as o3d
import os
import sys
import yaml

def pcd_to_pgm_pcd(input_pcd, output_pgm, output_yaml, resolution=0.1):
    """
    将PCD文件转换为PGM地图和YAML配置文件
    
    Args:
        input_pcd (str): 输入PCD文件路径
        output_pgm (str): 输出PGM文件路径
        output_yaml (str): 输出YAML配置文件路径
        resolution (float): 地图分辨率（米/像素）
    """
    # 读取PCD文件
    pcd = o3d.io.read_point_cloud(input_pcd)
    if len(pcd.points) == 0:
        raise ValueError("PCD文件中没有点云数据")
    
    # 将点云转换为numpy数组
    points = np.asarray(pcd.points)
    
    # 获取点云边界
    min_bound = np.min(points, axis=0)
    max_bound = np.max(points, axis=0)
    
    # 计算地图尺寸
    map_width = int((max_bound[0] - min_bound[0]) / resolution) + 20  # 添加边距
    map_height = int((max_bound[1] - min_bound[1]) / resolution) + 20  # 添加边距
    
    # 创建2D占用网格
    grid = np.zeros((map_height, map_width), dtype=np.uint8)
    
    # 将点云投影到2D网格
    for point in points:
        x = int((point[0] - min_bound[0]) / resolution) + 10  # 添加边距
        y = int((point[1] - min_bound[1]) / resolution) + 10  # 添加边距
        
        # 确保坐标在网格范围内
        if 0 <= x < map_width and 0 <= y < map_height:
            grid[y, x] = 255  # 标记为占用
    
    # 保存为PGM文件
    with open(output_pgm, 'wb') as f:
        f.write(b'P5\n')
        f.write(f'{map_width} {map_height}\n'.encode())
        f.write(b'255\n')
        grid.tofile(f)
    
    # 创建YAML配置文件
    yaml_data = {
        'image': os.path.basename(output_pgm),
        'resolution': resolution,
        'origin': [float(min_bound[0]) - 1.0, float(min_bound[1]) - 1.0, 0.0],
        'negate': 0,
        'occupied_thresh': 0.65,
        'free_thresh': 0.196
    }
    
    with open(output_yaml, 'w') as f:
        yaml.dump(yaml_data, f, default_flow_style=False)
    
    print(f"成功生成PGM地图: {output_pgm}")
    print(f"成功生成YAML配置: {output_yaml}")

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("用法: python3 pcd_to_map.py <input_pcd> <output_pgm> <output_yaml> [resolution]")
        print("示例: python3 pcd_to_map.py scans.pcd scans.pgm scans.yaml 0.1")
        sys.exit(1)
    
    input_pcd = sys.argv[1]
    output_pgm = sys.argv[2]
    output_yaml = sys.argv[3]
    resolution = float(sys.argv[4]) if len(sys.argv) > 4 else 0.1
    
    try:
        pcd_to_pgm_pcd(input_pcd, output_pgm, output_yaml, resolution)
    except Exception as e:
        print(f"转换失败: {e}")
        sys.exit(1)